package com.ndbg.demo.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class TreeUtil {
    public static <T extends MyTreeNode> List<T> buildTree(List<T> nodes) {
        // 使用 HashMap 存储所有节点
        Map<Long, T> nodeMap = new HashMap<>();
        for (T node : nodes) {
            nodeMap.put(node.getId(), node);
        }
        // 找到根节点，即没有父节点的节点
        List<T> rootNodes = new ArrayList<>();
        for (T node : nodes) {
            if (node.getParentId() == null || !nodeMap.containsKey(node.getParentId())) {
                rootNodes.add(node);
            }
        }
        // 递归生成树形结构
        for (T rootNode : rootNodes) {
            buildTreeRecursively(rootNode, nodeMap);
        }
        return rootNodes;
    }

    private static <T extends MyTreeNode> void buildTreeRecursively(T node, Map<Long, T> nodeMap) {
        // 找到所有子节点
        List<T> childNodes = new ArrayList<>();
        for (T childNode : nodeMap.values()) {
            if (childNode.getParentId() != null && childNode.getParentId().equals(node.getId())) {
                childNodes.add(childNode);
            }
        }
        // 递归生成子树
        for (T childNode : childNodes) {
            buildTreeRecursively(childNode, nodeMap);
        }
        // 将子节点添加到父节点中
        if (!childNodes.isEmpty()) {
            node.setChildren(childNodes);
        }
    }


    public interface MyTreeNode {
        Long getId();
        Long getParentId();
        void setChildren(List<? extends MyTreeNode> children);
    }
}
